import os
from ultralytics import YOLO

if __name__ == '__main__':
    # 读取模型
    model = YOLO(r'D:\pythonProject\autoDl-weed-detect\改进网络训练\weights\best.pt')

    # 指定输入和输出文件夹
    input_folder = r"D:\pythonProject\weed_detaction\datasets\predict_data"
    # output_folder = r"D:\pythonProject\weed_detaction\datasets\predicted_data"
    #
    # # 确保输出文件夹存在，否则创建它
    # if not os.path.exists(output_folder):
    #     os.makedirs(output_folder)

    # # 遍历输入文件夹中的所有图片
    # for filename in os.listdir(input_folder):
    #     if filename.endswith(('.jpg', '.jpeg', '.png', '.bmp')):  # 根据你的图片格式调整此条件
    #         input_filepath = os.path.join(input_folder, filename)
    #         output_filepath = os.path.join(output_folder, filename)  # 保持原始文件名
    #
    #         # 对单个图片进行预测并保存
    #         metrics = model.predict([input_filepath], save=True, output=output_filepath)
    metrics = model.predict(source=input_folder,
                            save=True,
                            device="cpu")

    # 如果你想要手动处理每个预测结果，可以按需添加相应逻辑
    # for m in metrics:
    #     box = m.boxes
    #     xywh = box.xywh
    #     cls = box.cls
    #     print(box, xywh, cls)

    # print("所有图片预测完成，并已保存至:", output_folder)